// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: xguo

#include "coin2/exchange/binance_base/api_util/rate_tracker.h"

#include <algorithm>
#include <sstream>
#include <string>

#include <boost/algorithm/string/replace.hpp>
#include <glog/logging.h>

#include "coin2/base/proto_util.h"
#include "coin2/exchange/base/order/rate_config.h"

namespace coin2::exchange::binance::api_util {

constexpr auto EXAMPLE_SPOT_BINANCE = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke18/binance/trade_key.json",
  "rate_limit_info": {
    "spot_binance": {
      "req_all": {
        "secs": 60,
        "reqs": 1200
      },
      "ord_sec": {
        "secs": 10,
        "reqs": 100
      },
      "ord_day": {
        "secs": 86400,
        "reqs": 200000
      }
    }
  }
},)";

constexpr auto EXAMPLE_FUTURES_BINANCE = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke18/binance/fut_trade_key.json",
  "rate_limit_info": {
    "futures_binance": {
      "req_all": {
        "secs": 60,
        "reqs": 2400
      },
      "ord_sec": {
        "secs": 10,
        "reqs": 300
      },
      "ord_min": {
        "secs": 60,
        "reqs": 1200
      }
    }
  }
},)";

constexpr auto EXAMPLE_FUTURES_BINANCE_DELIVERY = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke18/binance/fut_trade_key.json",
  "rate_limit_info": {
    "futures_binance": {
      "req_all": {
        "secs": 60,
        "reqs": 2400
      },
      "ord_min": {
        "secs": 60,
        "reqs": 1200
      }
    }
  }
},)";

template <class RateTracker>
void CheckFail(RateTracker&& rate_tracker, bool check_result) {
  if (!check_result) {
    LOG(INFO) << "Hit rate limit on [" << rate_tracker->GetName() << "]";
  }
}

BinanceRateTracker::BinanceRateTracker(
    const MarketExchangeApi& mea,
    const coin2::exchange::order::OrderConnectionConfig& cfg)
    : mea_{mea} {
  using RateConfig = coin2::exchange::base::order::RateConfig;
  if (mea.market == coin::proto::MarketType::Futures) {
    auto env_var = getenv("WHITELISTED_IP");
    whitelisted_ = (env_var && std::string(env_var) == std::string("1"));
    LOG(INFO) << "binance whitelisted ip is set: " << whitelisted_;
    coin2::exchange::order::RateLimitInfoFuturesBinance rate_config;

    // For backward compatibility.
    coin2::exchange::order::OrderConnectionConfig conn_cfg;
    if (whitelisted_) {
      auto* rl_info = conn_cfg.mutable_rate_limit_info()->mutable_futures_binance();
      rl_info->mutable_req_all()->set_secs(60);
      rl_info->mutable_req_all()->set_reqs(20000);
      rl_info->mutable_ord_sec()->set_secs(10);
      rl_info->mutable_ord_sec()->set_reqs(6666);
      rl_info->mutable_ord_min()->set_secs(60);
      rl_info->mutable_ord_min()->set_reqs(20000);
    }
    if (mea.api == "v1-delivery") {
      // for v1-swap, no 10-seconds restriction.
      auto* rl_info = conn_cfg.mutable_rate_limit_info()->mutable_futures_binance();
      rl_info->mutable_ord_sec()->set_secs(10);
      rl_info->mutable_ord_sec()->set_reqs(100000000000);
    }
    conn_cfg.MergeFrom(cfg);

    RateConfig::Get()->GetConfig(mea, conn_cfg, &rate_config);
    InitRateTracker(rate_config, mea.String(), EXAMPLE_FUTURES_BINANCE);
    RetrieveRateLimitInfoByIp(rate_config, cfg);
  } else if (
      mea.market == coin::proto::MarketType::Spot ||
      mea.market == coin::proto::MarketType::Margin) {
    std::string mea_str(fmt::format("Spot.{}.{}", ExchangeTypeToString(mea.exchange), mea.api));
    auto rate_limit_mea = MarketExchangeApi::FromString(mea_str);
    coin2::exchange::order::RateLimitInfoFuturesBinance rate_config;
    coin2::exchange::order::OrderConnectionConfig conn_cfg;
    auto* rl_info = conn_cfg.mutable_rate_limit_info()->mutable_spot_binance();
    rl_info->mutable_ord_min()->set_secs(60);
    rl_info->mutable_ord_min()->set_reqs(10000000000);
    RateConfig::Get()->GetConfig(rate_limit_mea, cfg, &rate_config);
    InitRateTracker(rate_config, rate_limit_mea.String(), EXAMPLE_SPOT_BINANCE);
    RetrieveRateLimitInfoByIp(rate_config, cfg);
  } else {
    CHECK(false) << "unsupported mea: " << mea.String();
  }
}

void BinanceRateTracker::RetrieveRateLimitInfoByIp(
    const google::protobuf::Message& default_cfg,
    const coin2::exchange::order::OrderConnectionConfig& conn_cfg) {
  // retrieve default rate limit config
  coin2::exchange::order::RateLimitInfoFuturesBinance default_rate_config;
  default_rate_config.MergeFrom(default_cfg);
  rate_limit_by_ip_cfg_.emplace("", std::move(default_rate_config));
  // retrieve bind ip rate limit config
  for (const auto& bind_ip_conn : conn_cfg.bind_ip_conns()) {
    if (!bind_ip_conn.has_bind_ip() || bind_ip_conn.bind_ip().empty()) {
      continue;
    }
    if (!bind_ip_conn.has_rate_limit_info()) {
      continue;
    }
    coin2::exchange::order::RateLimitInfoFuturesBinance rate_config;
    using RateConfig = coin2::exchange::base::order::RateConfig;
    RateConfig::MergeFromConfig(
        RateConfig::GetMeStringFromMea(mea_),
        bind_ip_conn.rate_limit_info(),
        &rate_config);
    rate_limit_by_ip_cfg_.emplace(bind_ip_conn.bind_ip(), std::move(rate_config));
  }
}

auto BinanceRateTracker::GetReqRateTracker(const std::string& ip) -> RateTrackerType& {
  // get rate tracker by ip
  if (req_rate_tracker_.count(ip) == 0) {
    const auto& rate_config = rate_limit_by_ip_cfg_.count(ip) == 0 ? rate_limit_by_ip_cfg_[""]
                                                                   : rate_limit_by_ip_cfg_[ip];
    req_rate_tracker_[ip] =
        std::make_unique<FixedWindowRateTracker>(rate_config.req_all().secs() * 1'000'000'000LL);
    // according to doc, max points is 100, but test shows it is 7000.
    req_rate_tracker_[ip]->SetRateLimitInfo(
        rate_config.req_all().reqs(),
        rate_config.req_all().reqs());
    req_rate_tracker_[ip]->SetName("binance ip rate tracker(1 min): " + ip);
  }
  return req_rate_tracker_[ip];
}

void BinanceRateTracker::InitRateTracker(
    const RateLimitInfoFuturesBinance& rate_config,
    std::string_view mea_str,
    std::string_view example) {
  if (!rate_config.has_req_all() || !rate_config.has_ord_min() || !rate_config.has_ord_day() ||
      !rate_config.has_ord_sec() || !rate_config.has_waf()) {
    coin2::exchange::base::order::RateConfig::CheckFail(mea_str, example);
  }
  const auto& req_all = rate_config.req_all();
  const auto& ord_min = rate_config.ord_min();
  const auto& ord_day = rate_config.ord_day();
  const auto& ord_sec = rate_config.ord_sec();
  const auto& waf = rate_config.waf();
  LOG(INFO) << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
            << "\nOG request rate was set as below:\n"
            << "\nreq_all: [" << req_all.reqs() << " (reqs) / " << req_all.secs() << " (seconds)]"
            << "\nord_sec: [" << ord_sec.reqs() << " (reqs) / " << ord_sec.secs() << " (seconds)]"
            << "\nord_min: [" << ord_min.reqs() << " (reqs) / " << ord_min.secs() << " (seconds)]"
            << "\nord_day: [" << ord_day.reqs() << " (reqs) / " << ord_day.secs() << " (seconds)]"
            << "\nwaf: [" << waf.reqs() << " (reqs) / " << waf.secs() << " (seconds)]"
            << "\n\nOr you can configure customized rate in driver config. For example: " << example
            << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~";

  constexpr int maxn = std::numeric_limits<int>::max();

  order_sec_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(ord_sec.secs() * 1'000'000'000LL);
  if (ord_sec.reqs() == 0) {
    order_sec_rate_tracker_->SetRateLimitInfo(maxn, maxn);
  } else {
    order_sec_rate_tracker_->SetRateLimitInfo(ord_sec.reqs(), ord_sec.reqs());
  }
  order_sec_rate_tracker_->SetName("binance_order_rate_tracker(10 sec)");

  order_min_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(ord_min.secs() * 1'000'000'000LL);
  if (ord_min.reqs() == 0) {
    order_min_rate_tracker_->SetRateLimitInfo(maxn, maxn);
  } else {
    order_min_rate_tracker_->SetRateLimitInfo(ord_min.reqs(), ord_min.reqs());
  }
  order_min_rate_tracker_->SetName("binance_order_rate_tracker(1 min)");

  order_day_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(ord_day.secs() * 1'000'000'000LL);
  if (ord_day.reqs() == 0) {
    order_day_rate_tracker_->SetRateLimitInfo(maxn, maxn);
  } else {
    order_day_rate_tracker_->SetRateLimitInfo(ord_day.reqs(), ord_day.reqs());
  }
  order_day_rate_tracker_->SetName("binance_order_rate_tracker(1 day)");

  waf_rate_tracker_ = std::make_unique<FixedWindowRateTracker>(waf.secs() * 1'000'000'000LL);
  if (waf.reqs() == 0) {
    waf_rate_tracker_->SetRateLimitInfo(maxn, maxn);
  } else {
    waf_rate_tracker_->SetRateLimitInfo(waf.reqs(), waf.reqs());
  }
  waf_rate_tracker_->SetName("binance_waf_rate_tracker(5 min)");

  // rate limit report
  rate_limit_report_writer_ = std::make_unique<RateLimitReportWriter>();
  rate_limit_report_writer_->SetRateLimitInfoString(&rate_config);
}

bool BinanceRateTracker::HasEnoughRateLimit(int64_t timestamp, ActionType action_type) {
  if (CheckEnableTime(timestamp, action_type) && CheckRateLimit(timestamp, action_type)) {
    return true;
  }

  rate_limit_report_writer_->SetHitRateLimit(timestamp, action_type);
  return false;
}

// Only called in native private client. SetBindIp must be called before this function.
void BinanceRateTracker::Push(int32_t num_reqs, int64_t timestamp, ActionType action_type) {
  waf_rate_tracker_->Push(1, timestamp);
  if (action_type == ActionType::PLACE) {
    if (mea_.market == MarketType::Spot) {
      GetReqRateTracker(bind_ip_)->Push(num_reqs, timestamp);
    }
    order_day_rate_tracker_->Push(num_reqs, timestamp);
    order_sec_rate_tracker_->Push(num_reqs, timestamp);
    order_min_rate_tracker_->Push(num_reqs, timestamp);
  } else {
    GetReqRateTracker(bind_ip_)->Push(num_reqs, timestamp);
  }
}

// Called in native private client. SetBindIp must be called.
void BinanceRateTracker::UpdateReqRateInfo(const std::string& channel, void* data) {
  auto info = nlohmann::json::parse(*reinterpret_cast<std::string*>(data));

  nlohmann::json json;
  json["update"] = info;

  bind_ip_ = info["bind_ip"].get<std::string>();

  // enable time
  int64_t retry_after = info["retry-after"].get<int>();
  if (retry_after > 0) {
    SetEnableTime((retry_after + 2) * 1'000'000'000LL + GetCurrentTimestamp());
  }

  // IP specific
  int req_rate_limit_consumed = info["x-mbx-used-weight-1m"].get<int>();
  if (req_rate_limit_consumed >= 0) {
    auto rate_limit = GetReqRateTracker(bind_ip_)->MaxReq();
    auto rate_limit_remaining = rate_limit - req_rate_limit_consumed;
    GetReqRateTracker(bind_ip_)->SetRateLimitInfo(rate_limit, rate_limit_remaining);
    int64_t arr[] = {rate_limit, rate_limit_remaining};
    json["req_rate_limit"] = arr;
  }

  // Only apply to spot trading, for futures, maxreq is set to a very big number
  int order_day_rate_limit_consumed = info["x-mbx-order-count-1d"].get<int>();
  if (order_day_rate_limit_consumed >= 0) {
    auto rate_limit = order_day_rate_tracker_->MaxReq();
    auto rate_limit_remaining = rate_limit - order_day_rate_limit_consumed;
    order_day_rate_tracker_->SetRateLimitInfo(rate_limit, rate_limit_remaining);
    json["order_day_rate_limit"] = {rate_limit, rate_limit_remaining};
  }

  // For Spot and fapi, for dapi, maxreq is set to a big number.
  int order_sec_rate_limit_consumed = info["x-mbx-order-count-10s"].get<int>();
  if (order_sec_rate_limit_consumed >= 0) {
    auto rate_limit = order_sec_rate_tracker_->MaxReq();
    auto rate_limit_remaining = rate_limit - order_sec_rate_limit_consumed;
    order_sec_rate_tracker_->SetRateLimitInfo(rate_limit, rate_limit_remaining);
    json["order_sec_rate_limit"] = {rate_limit, rate_limit_remaining};
  }

  // Only for fapi and dapi, for spot, maxreq is set to a big number.
  int order_min_rate_limit_consumed = info["x-mbx-order-count-1m"].get<int>();
  if (order_min_rate_limit_consumed >= 0) {
    auto rate_limit = order_min_rate_tracker_->MaxReq();
    auto rate_limit_remaining = rate_limit - order_min_rate_limit_consumed;
    order_min_rate_tracker_->SetRateLimitInfo(rate_limit, rate_limit_remaining);
    json["order_min_rate_limit"] = {rate_limit, rate_limit_remaining};
  }

  LOG_EVERY_NS(INFO, 5'000'000'000, GetCurrentTimestamp()) << "update rate limit: \n"
                                                           << json.dump();
}

void BinanceRateTracker::SetRateLimitInfoString(const google::protobuf::Message* message) {
  if (!rate_limit_report_writer_) {
    return;
  }
  rate_limit_report_writer_->SetRateLimitInfoString(message);
}

void BinanceRateTracker::WriteRateLimitReportMessage(OrderLoggerCB callback) {
  if (!callback) {
    return;
  }
  rate_limit_report_writer_->Write(callback);
}

double BinanceRateTracker::GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) {
  double waf_ratio = 0.0;
  {
    auto max_reqs = waf_rate_tracker_->MaxReq();
    auto remaining = waf_rate_tracker_->GetRemaining(timestamp);
    waf_ratio = 1.0 * (max_reqs - remaining) / max_reqs;
  }

  switch (action_type) {
    case ActionType::PLACE:
    case ActionType::AMEND: {
      double req_ratio = 0.;
      {
        auto max_reqs = GetReqRateTracker("")->MaxReq();
        auto remaining = GetReqRateTracker("")->GetRemaining(timestamp);
        req_ratio = 1.0 * (max_reqs - remaining) / max_reqs;
      }

      double min_ratio = 0.;
      {
        auto max_reqs = order_min_rate_tracker_->MaxReq();
        auto remaining = order_min_rate_tracker_->GetRemaining(timestamp);
        min_ratio = 1.0 * (max_reqs - remaining) / max_reqs;
      }

      double day_ratio = 0.;
      {
        auto max_reqs = order_day_rate_tracker_->MaxReq();
        auto remaining = order_day_rate_tracker_->GetRemaining(timestamp);
        day_ratio = 1.0 * (max_reqs - remaining) / max_reqs;
      }

      double sec_ratio = 0.;
      {
        auto max_reqs = order_sec_rate_tracker_->MaxReq();
        auto remaining = order_sec_rate_tracker_->GetRemaining(timestamp);
        sec_ratio = 1.0 * (max_reqs - remaining) / max_reqs;
      }
      if (mea_.market == MarketType::Spot) {
        return std::max({req_ratio, min_ratio, sec_ratio, day_ratio, waf_ratio});
      } else {
        return std::max({min_ratio, sec_ratio, day_ratio, waf_ratio});
      }
    }
    case ActionType::CANCEL: {
      auto max_reqs = GetReqRateTracker("")->MaxReq();
      auto remaining = GetReqRateTracker("")->GetRemaining(timestamp);
      return std::max(1.0 * (max_reqs - remaining) / max_reqs, waf_ratio);
    }
    case ActionType::QUERY: {
      // always return 0
      auto max_reqs = GetReqRateTracker(bind_ip_)->MaxReq();
      auto remaining = GetReqRateTracker(bind_ip_)->GetRemaining(timestamp);
      return 0.0 * (max_reqs - remaining) / max_reqs;
    }
    default: {
      // shouldn't reached here
      LOG(ERROR) << "ops~ something is wrong, action_type=" << action_type;
      return 0.;
    }
  }
}

bool BinanceRateTracker::CheckEnableTime(int64_t timestamp, ActionType action_type) {
  if (enable_time_by_ip_map_.count(bind_ip_) == 0) {
    return true;
  }

  auto enable_time = enable_time_by_ip_map_[bind_ip_];

  if (timestamp >= enable_time) {
    return true;
  }

  LOG_EVERY_N(WARNING, 50) << "Hit api temporary disabled until " << enable_time << ", "
                           << action_type;
  return false;
}
bool BinanceRateTracker::CheckRateLimit(int64_t timestamp, ActionType action_type) {
  bool ret = true;

  ret &= waf_rate_tracker_->GetRemaining(timestamp) > 100;
  CheckFail(waf_rate_tracker_, ret);
  if (action_type == ActionType::PLACE) {
    if (mea_.market == MarketType::Spot) {
      ret &= GetReqRateTracker("")->GetRemaining(timestamp) > 20;
      CheckFail(GetReqRateTracker(""), ret);
    }

    ret &= order_day_rate_tracker_->GetRemaining(timestamp) > 20;
    CheckFail(order_day_rate_tracker_, ret);

    ret &= order_min_rate_tracker_->GetRemaining(timestamp) > 5;
    CheckFail(order_min_rate_tracker_, ret);

    ret &= order_sec_rate_tracker_->GetRemaining(timestamp) > 5;
    CheckFail(order_sec_rate_tracker_, ret);
  } else if (action_type == ActionType::CANCEL) {
    ret &= GetReqRateTracker("")->GetRemaining(timestamp) > 20;
    CheckFail(GetReqRateTracker(""), ret);
  } else {
    // Only used in native client
    if (bind_ip_.empty()) {
      ret &= GetReqRateTracker("")->GetRemaining(timestamp) > GetReqRateTracker("")->MaxReq() / 2;
      CheckFail(GetReqRateTracker(bind_ip_), ret);
    } else {
      ret &= GetReqRateTracker(bind_ip_)->GetRemaining(timestamp) > 20;
      CheckFail(GetReqRateTracker(bind_ip_), ret);
    }
  }

  return ret;
}

}  // namespace coin2::exchange::binance::api_util
